'''
Created on May 10, 2011

5/10/11 Replaced QUEUE directive with FORWARD

@author: yev
'''

import array
from mpi4py import MPI
import xml.etree.cElementTree as cET
import os
import Queue
from random import randint, choice
import string
import sys
import tempfile
import time
from threading import Thread
import World 
import collections
import numpy
import Log

PRIVATE_QUEUE = {}
RESULTS_QUEUES = {}
luigi_logger = Log.create_logger("Wario3")

PIPE_COUNTER = 0

NEW_LINE = "\n"

BUFFER_SLEEP = True # puts thread to sleep if buffer is full

def str2bool(v):
    if v is None:
        return False
    return v.lower() in ["yes", "true", "t", "1"]


# Ranks where data is coming from
def find_ranks_for_task(rank_to_task, task_of_interest):
    source_ranks = []
    for rank in rank_to_task:
        if rank_to_task[rank] == task_of_interest:
            source_ranks.append(rank)
    return source_ranks

# Ranks where data is going
def find_ranks_for_tasks(rank_to_task, tasks_of_interest):
    destination_ranks = []
    for rank in rank_to_task:
        if rank_to_task[rank] in tasks_of_interest:
            destination_ranks.append(rank)
    return destination_ranks

def work(xml_file):
    PIPE_COUNTER = 0
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    name = MPI.Get_processor_name()
    # sending id
    comm.send(name, dest=0, tag=World.WHO)
    rank_to_task = {}
    rank_to_task = comm.recv(rank_to_task, source=0, tag=World.ASSIGNMENT)
    if not rank_to_task:
        return
    
    assigned_task = rank_to_task[rank]
    
    
    luigi_logger.debug ("I am Luigi %d on %s. My assigned task is %s. My PID is %d" %(rank, name, assigned_task, os.getpid()))
    pipes = {}
    commands = []
    global BUFFER_SLEEP
    for event, elem in cET.iterparse(xml_file, events = ("start", "end")):
        if event == "end" and "task" in elem.tag:
            print elem.attrib["id"]
            if elem.attrib["id"] == assigned_task:
    
                print elem.attrib["weight"]
                if elem.attrib["weight"] == "1":
                    BUFFER_SLEEP = False  # There is no point in pausing the thread if it is a singleton
                              
                    print ("Luigi %d will not be put to sleep" %(rank))
                else:
                    BUFFER_SLEEP = True
               
                for item in elem.getiterator("send"):
                    if item.text is None:
                        send = None
                        luigi_logger.debug ("Luigi %d was not assigned a send task" %(rank))
                    else:
                        destinations = item.text.split(",")
                        for destination in destinations:    # assign a list to every task that will receive data via MPI
                            RESULTS_QUEUES[int(destination)] = {}
                        receivers = find_ranks_for_tasks(rank_to_task, destinations)
                        offset = item.attrib["offset"]
                        send = Send(int(assigned_task), offset, receivers, rank_to_task)
    
                for item in elem.getiterator("receive"):
                    if item.text is None:
                        receive = None
                    else:
                        source = int(item.text)
                        senders = find_ranks_for_task(rank_to_task, str(source))
                        receive = Receive(source,senders)
    
                for pipe in elem.getiterator("pipe"):
                    id = pipe.attrib["id"]
                    incoming = str2bool(pipe.attrib["incoming"])
                    if incoming:
                        pipe_index = int(pipe.attrib["index"])
                        PRIVATE_QUEUE[PIPE_COUNTER] = Queue.Queue()
                        pipe = Pipe(incoming, private_queue_index = PIPE_COUNTER, incoming_index = pipe_index)
                        PIPE_COUNTER += 1
                    else:    
                        if pipe.attrib.get("path") != None: # pipe path or file path is specified
                            pipe = Pipe(incoming, set_path = pipe.attrib["path"])
                        else:
                            pipe = Pipe(incoming)

                    pipes[id] = pipe
    
                for command_desc in elem.getiterator("command"):
                    command_text = command_desc.text.split()
                    translated_command = ''
                    for word in command_text:
                        pipe_index = word.find('%')
                        if pipe_index >= 0:
                            word = word.strip('%')
                            word = pipes[word].getPath()
                        translated_command += word
                        translated_command += ' '
                    
                    out_hint = command_desc.attrib["out"]
                    pipe_index = out_hint.find('%')
                    if pipe_index >= 0:                      # internal pipe
                        out_hint = out_hint.strip('%')
                        std_out = pipes[out_hint].getPath()
                    elif out_hint == "NONE":
                        std_out = None
                    forward_stdout = str2bool(command_desc.attrib["forward"])
                    if forward_stdout is True:
                        index = command_desc.attrib["index"]
                        index = int(index)
                        for key in RESULTS_QUEUES:
                            RESULTS_QUEUES[key][index] = collections.deque()
                        command = Command(translated_command, std_out, forward_stdout, index)
                    else:
                        command = Command(translated_command, std_out, forward_stdout)
                    commands.append(command)
                    
                elem.clear() 
        
    threads = []
    for command in commands:
        PIPE_COUNTER += 1
        thread = RunExternalCommand(command, PIPE_COUNTER)
        thread.setName(command.getCommand())
        thread.start()
        threads.append(thread)
    
    if send is None and receive is not None:
        luigi_logger.debug ("Luigi %d is in RECEIVING mode" %(rank))
        do_receive(receive, pipes, rank)
    elif send is not None and receive is None:
        luigi_logger.debug ("Luigi %d is in SENDING mode" %(rank))
        do_send(send, threads)
    elif send != None and receive != None:
        luigi_logger.debug ("Luigi %d is in RECEIVING AND SENDING mode" %(rank))
        do_receive_and_send(receive, pipes, rank, send, threads)
    
    for thread in threads:
        thread.join()


class Command:
    def __init__(self, command, out, forward_stdout, index = 0):
        self.command = command                  # command itself
        self.stdOut = out                       # place to put std out if any
        self.forward_stdout = forward_stdout    # whether to forward the output
        self.index = index
    
    def getCommand(self):
        commandList = []
        for word in self.command.split():
            result = ""
            result += word
            commandList.append(result)
        return commandList
    
    def getStdOut(self):
        return self.stdOut
    
    def get_forward_stdout(self):
        return self.forward_stdout

class Send:
    def __init__(self, destination, offset, receivers, rank_to_task_mapping):
        self.destination = destination
        self.offset = offset
        self.receivers = receivers
        self.rank_to_task_mapping = rank_to_task_mapping
    
    def get_rank_to_task_mapping(self):
        return self.rank_to_task_mapping

class Receive:
    def __init__(self, source, senders):
        self.source = source
        self.senders = senders


class Internal_Pipe:
    def __init__(self, queue_index, path):
        self.queue_index = queue_index
        self.path = path
        self.worker = PipeWorker(self.path, self.queue_index)
        self.worker.daemon = True
        self.worker.start()
    
    def writeData(self, data):
        PRIVATE_QUEUE[self.queue_index].put(data)
    
    def closeFile(self):
        self.worker.done = True
        PRIVATE_QUEUE[self.queue_index].join()
        self.worker.fh.close()



class Pipe:
    def __init__(self, incoming, private_queue_index = None, set_path = None, incoming_index = -1):
        self.incoming = incoming
        self.incoming_index = incoming_index
        if set_path is not None:
            self.path = set_path
            open(self.path, 'w').close()
        else:    
            self.path = tempfile.gettempdir() + "/" + randomString()
            os.mkfifo(self.path)
        
        if self.incoming == True:
            self.private_queue_index = private_queue_index    
            self.worker = PipeWorker(self.path, self.private_queue_index)
            self.worker.daemon = True
            self.worker.start()
             
    def writeData(self, data):
        PRIVATE_QUEUE[self.private_queue_index].put(data)
        return PRIVATE_QUEUE[self.private_queue_index].qsize()
    
    def closeFile(self):
        if self.incoming == True:
            self.worker.done = True
            PRIVATE_QUEUE[self.private_queue_index].join()
        return
    
    def getQueueSize(self):
        return PRIVATE_QUEUE[self.private_queue_index].qsize()
    
    def getPath(self):
        return self.path
            
    def __str__(self):
        return ("Hello" + self.path)
 
class PipeWorker(Thread):
    def __init__(self, path, private_queue_index):
        Thread.__init__(self)
        self.pipePath = path
        self.number = 0
        self.done = False
        self.private_queue_index = private_queue_index
        self.fh = None
    def run(self):
        time.sleep(1)
        self.fh = open(self.pipePath, 'w')
        while True:
            if self.done is True and PRIVATE_QUEUE[self.private_queue_index].empty():
                break                
            
            if not  PRIVATE_QUEUE[self.private_queue_index].empty():
                work = PRIVATE_QUEUE[self.private_queue_index].get()
            else:
                time.sleep(1)
                continue
            
            if isinstance(work, str):
                self.fh.write(work)
            else:
                self.fh.write(work.tostring())
            self.number += 1   
            PRIVATE_QUEUE[self.private_queue_index].task_done()
        self.fh.close()
                           
        
class RunExternalCommand(Thread):
    def __init__(self, externalCommand, PIPE_COUNTER):
        Thread.__init__(self)
        self.externalCommand = externalCommand
        self.pipe_counter = PIPE_COUNTER
        PRIVATE_QUEUE[self.pipe_counter] = Queue.Queue(0)
        self.command = ""
        for word in externalCommand.getCommand():
            self.command += word
            self.command += " "

        if self.externalCommand.getStdOut() is None:
            self.output = None
            self.proc = os.popen(self.command)
        else:
            self.output = self.externalCommand.getStdOut()
            self.proc = os.popen(self.command)

    def is_text(self, text):
        text_characters = "".join(map(chr, range(32, 127)) + list("\n\r\t\b"))
        _null_trans = string.maketrans("", "")
        if "\0" in text:
            return False
    
        if not text:  # Empty files are considered text
            return 1
        # Get the non-text characters (maps a character to itself then
        # use the 'remove' option to get rid of the text characters.)
        t = text.translate(_null_trans, text_characters)
        # If more than 30% non-text characters, then
        # this is considered a binary file
        if float(len(t)) / float(len(text)) > 0.30:
            return False
        return True
            
    def run(self):
        checked_for_type = False
        is_text = False        
        if self.output == None:       # Nothing is being written to file
            luigi_logger.debug ("Thread is STARTING: command=%s, capture=%s" %(self.command, self.externalCommand.getStdOut()))
            if self.externalCommand.get_forward_stdout(): # Sending stdout via MPI    
                while True:
                    if not checked_for_type:
                        line = self.proc.readline()
                        is_text = self.is_text(line)
                        if not is_text:
                            luigi_logger.debug ("Determined that command %s is producing binary output" %(self.command))
                        checked_for_type = True
                    elif not is_text:
                        line = self.proc.read(65536)
                    else:
                        line = self.proc.readline()
                    for task in RESULTS_QUEUES:
                        for index in RESULTS_QUEUES[task]:
                            if index == self.externalCommand.index:
                                RESULTS_QUEUES[task][index].appendleft(line)
#                                print self.command, task, index, line
                        
                    sys.stdout.flush()
                    if not line:
                        try:
                            self.proc.close()
                        except IOError:
                            pass
                        luigi_logger.debug ("Thread is Exiting: command=%s, capture=%s" %(self.command, self.externalCommand.getStdOut()))
                        break
            else:                                       # Stay with the process until it exits
                exit_code = self.proc.close()
               
        else:
            luigi_logger.debug ("Thread is STARTING: command=%s, capture=%s" %(self.command, self.externalCommand.getStdOut()))
            fh = open(self.externalCommand.getStdOut(), 'w')    # Writing to a file object
            send = False
            send = self.externalCommand.get_forward_stdout()
            while True:
                if not checked_for_type:
                    line = self.proc.readline()
                    is_text = self.is_text(line)
                    if not is_text:
                        print ("Determined that command %s is producing binary output" %(self.command))
                    checked_for_type = True
                elif not is_text:
                    line = self.proc.read(65536)
                else:
                    line = self.proc.readline()
                sys.stdout.flush()
                fh.write(line)
                if send:
                    for task in RESULTS_QUEUES:
                        for index in RESULTS_QUEUES[task]:
                            if index == self.externalCommand.index:
                                RESULTS_QUEUES[task][index].appendleft(line)
#                                print self.command, task, index, line
                if not line:
                    fh.close()
                    try:
                        self.proc.close()
                    except IOError:
                        pass
                    
                    luigi_logger.debug ("Thread is Exiting: command=%s, capture=%s" %(self.command, self.externalCommand.getStdOut()))
                    break

def randomString():
    chars = string.ascii_letters
    return "".join(choice(chars) for x in range(randint(6, 6)))

def prepare_message(offset, threads, task):
    """
    Form message to send via MPI
    
    Keyword arguments:
    offset -- message size (typically number of lines)
    threads -- threads running; used to check if it is time to call it quits
    task -- the task id requesting data
    return -- dictionary (command index -> message, eof)
    """
    start_time = time.time()
    data = []
    transmission = {}
    eof = False
    for index in RESULTS_QUEUES[task]:
        retrieved = 0
        data = []
        luigi_logger.debug("Preparing message for task %d index %d" %(task, index))
        while retrieved < offset:
            if not are_threads_alive(threads):
                if len(RESULTS_QUEUES[task][index]) == 0:
                    eof = True
                    break
            try:            
                bleeding_edge = RESULTS_QUEUES[task][index].pop()
#                print task, index, bleeding_edge
            except IndexError: # The dequeue is empty
                luigi_logger.debug("Sleeping on message for task %d index %d" %(task, index))
                time.sleep(0.5)
                continue
            data.append(bleeding_edge)
            retrieved += 1
        message = "".join(data)
        transmission[index] = message
#        print message
    luigi_logger.debug("Message for task %d was prepared in %f seconds" %(task, time.time() - start_time))
    return [transmission, eof]
    
def do_send(send, threads):
    dummy = array.array('i', range(2))
    status = MPI.Status()
    comm = MPI.COMM_WORLD
    offset = int(send.offset)
    rank_to_task = send.get_rank_to_task_mapping()
    req = None
    receivers = send.receivers
            
    while len(receivers) > 0:
        for rank in receivers:
            if not comm.Iprobe(source=rank, tag=MPI.ANY_TAG, status=status): # is there anyone waiting?
#                time.sleep(0.1)
                continue
            task = int(rank_to_task[rank])
            transmission, eof = prepare_message(offset, threads, task)
            start_time = time.time()
            comm.Recv([dummy, 2, MPI.INT], source=status.source, tag=send.destination) # get the rank of the waiting process
            luigi_logger.debug ("Attempting to send %d to %d for task %d" %(len(transmission), status.source, task))
#            print "Transmission keys: ", sorted(transmission.iterkeys())
            for index in sorted(transmission.iterkeys()):
                if req is not None: # Did the previous recipient receive the work?
                    req.Wait()
                message = transmission[index]
#                print message       
                req = comm.Isend([message, len(message), MPI.CHAR], dest=status.source, tag=int(index))
                luigi_logger.debug ("Successful sending to %d (%f), index %d" %(status.source, time.time() - start_time, int(index)))
            
            if eof:
                luigi_logger.debug ("Sender detected EOF for task %d" %(task))
                ranks_affected_by_eof = find_ranks_for_task(rank_to_task, rank_to_task[rank])
                if req is not None:
                    req.Wait()
#                luigi_logger.debug ("ranks affected: ", ranks_affected_by_eof)
                while len(ranks_affected_by_eof) > 0:
                    for affected_rank in ranks_affected_by_eof:
                        if not comm.Iprobe(source=affected_rank, tag=MPI.ANY_TAG, status=status):
                            continue
                        luigi_logger.debug ("Sending SEACREST_OUT to %d" %(affected_rank))
                        comm.Send([dummy, 2, MPI.INT], dest=affected_rank, tag=World.SEACREST_OUT)
                        comm.Recv([dummy, 2, MPI.INT], source=affected_rank, tag=World.SEACREST_OUT_ACK)
                        luigi_logger.debug ("Received SEACREST_OUT ACK from %d" %(rank))
                        ranks_affected_by_eof.remove(affected_rank)
                        receivers.remove(affected_rank)
                return
                
def do_receive(receive, pipes, luigi_rank):
    dummy = array.array('i', range(2))
    comm = MPI.COMM_WORLD
    status = MPI.Status()
    number_of_senders = len(receive.senders)
    
    data = numpy.empty(10, dtype = 'c')
    
    # Get indices of the incoming pipes
    indices = []
    for pipe in pipes:
        if pipes[pipe].incoming == True:
            indices.append(int(pipes[pipe].private_queue_index))

    while number_of_senders > 0:
        for rank in receive.senders:
#            print ("Asking %d" %(rank))
            comm.Send([dummy, 2, MPI.INT], dest=rank, tag=receive.source)
            while not comm.Iprobe(source=rank, tag=MPI.ANY_TAG, status=status):
                time.sleep(0.1)
            
            if status.tag == World.SEACREST_OUT:
                luigi_logger.debug ("Luigi %d is receiving that Seacrest is out from %d" %(luigi_rank, rank))
                comm.Recv([dummy, 2, MPI.INT], source=rank, tag=World.SEACREST_OUT)
                comm.Send([dummy, 2, MPI.INT], dest=rank, tag=World.SEACREST_OUT_ACK)
                receive.senders.remove(rank)
                luigi_logger.debug ("%d Remain" %(len(receive.senders)))
                number_of_senders -= 1
                if number_of_senders == 0:
                    luigi_logger.debug ("Number of Senders for Luigi %d = %d"%(luigi_rank, number_of_senders))
                    pipe_list = pipes.keys()
                    # pipes are closed in alphabetical order, the order of closures matters
                    pipe_list.sort()
                    for pipe in pipe_list:
                        if  pipes[pipe].incoming == True:
                            luigi_logger.debug ("Luigi %d closing pipe %s" %(luigi_rank, pipe))
                            pipes[pipe].closeFile()
                            luigi_logger.debug ("Luigi %d closed pipe %s" %(luigi_rank, pipe))
                    return
            else:
                luigi_logger.debug ("Luigi %d is receiving new work from rank %d with tag %d" %(luigi_rank, rank, status.tag ))               
                data = numpy.empty(status.count, dtype = 'c')
                comm.Recv([data, status.count, MPI.CHAR], source=rank, tag=status.tag)
#                print data.tostring()
                feed_incoming_pipes_wait(data, pipes, luigi_rank, status.tag)                            

def do_receive_and_send(receive, pipes, luigi_rank, send, threads):
    dummy = array.array('i', range(2))
    comm = MPI.COMM_WORLD
    luigi_name = MPI.Get_processor_name()
    status = MPI.Status()
    number_of_senders = len(receive.senders)
    rank_to_task = send.get_rank_to_task_mapping()
    receivers = send.receivers
    offset = int(send.offset)
    req = None
    data = numpy.empty(10, dtype = 'c')
    wrap_up = False
        
    while True:
        if number_of_senders > 0:
            for rank in receive.senders:
#                print "asked ", rank
                comm.Send([dummy, 2, MPI.INT], dest=rank, tag=receive.source)              
                while not comm.Iprobe(source=rank, tag=MPI.ANY_TAG, status=status):
                    time.sleep(0.1)

                if status.tag == World.SEACREST_OUT:
                    comm.Send([dummy, 2, MPI.INT], dest=rank, tag=World.SEACREST_OUT_ACK)
                    number_of_senders -= 1
                    receive.senders.remove(rank)
                    luigi_logger.debug ("Luigi %d is receiving SEACREST_OUT from %d" %(luigi_rank, rank))
                else:
                    luigi_logger.debug ("Luigi %d is receiving new work from %d" %(luigi_rank, rank))
                    data = numpy.empty(status.count, dtype = 'c')
                    comm.Recv([data, status.count, MPI.CHAR], source=rank, tag=status.tag)
#                    print data, status.tag
                    feed_incoming_pipes_wait(data, pipes, luigi_rank, status.tag)
                
        elif not wrap_up:       # no more data are to be received
            # pipes are closed in alphabetical order, the order of closures matters
            pipe_list = pipes.keys()
            pipe_list.sort()
            for pipe in pipe_list:
                if  pipes[pipe].incoming == True:
                    pipes[pipe].closeFile()
            wrap_up = True
            luigi_logger.debug ("Luigi %d wrapped up" %(luigi_rank))
                        
        for rank in receivers:
            task = int (rank_to_task[rank])
            if not comm.Iprobe(source=rank, tag=MPI.ANY_TAG, status=status):
                continue            
            
            empty_queue = True
            for index in RESULTS_QUEUES[task]:
                if len(RESULTS_QUEUES[task][index]) > 0 or are_threads_alive(threads) is False:
                    empty_queue = False
                    break
            if empty_queue is True:
                continue
            
            comm.Recv([dummy, 2, MPI.INT], source=rank, tag=send.destination) # get the rank of the waiting process
            luigi_logger.debug ("Luigi %d is requesting message for task % d" %(luigi_rank, task))
            transmission, eof = prepare_message(offset, threads, task)
            luigi_logger.debug ("Luigi %d is sending message for task % d" %(luigi_rank, task))
            for index in sorted(transmission.iterkeys()):
                if req is not None: # Did the previous recipient receive the message
                    req.Wait()
                message = transmission[index]
                req = comm.Isend([message, len(message), MPI.CHAR], dest=status.source, tag=index)
            if eof:
                luigi_logger.debug ("Receiver / Sender Luigi %d detected EOF for task %d" %(luigi_rank, task))
                ranks_affected_by_eof = find_ranks_for_task(rank_to_task, rank_to_task[rank])
#                luigi_logger.debug ("ranks affected: ", ranks_affected_by_eof)
                while len(ranks_affected_by_eof) > 0:
                    for affected_rank in ranks_affected_by_eof:
                        if not comm.Iprobe(source=affected_rank, tag=MPI.ANY_TAG, status=status):
                            continue
                        luigi_logger.debug ("Sending SEACREST_OUT to %d" %(affected_rank))
                        comm.Send([dummy, 2, MPI.INT], dest=affected_rank, tag=World.SEACREST_OUT)
                        comm.Recv([dummy, 2, MPI.INT], source=affected_rank, tag=World.SEACREST_OUT_ACK)
                        luigi_logger.debug ("Received SEACREST_OUT ACK from %d" %(rank))
                        ranks_affected_by_eof.remove(affected_rank)
                        receivers.remove(affected_rank)
        
        if wrap_up and len(receivers) == 0: # Exit when nothing to send or receive
            return

def data_in_queue():
    for queue in RESULTS_QUEUES.iterkeys():
        if len(RESULTS_QUEUES[queue]) > 0:
            return True
    return False

def are_threads_alive(threads):
    for thread in threads:
        if thread.is_alive():
            return True
    return False


def we_are_done(threads, dequeue):
    if not are_threads_alive(threads):
        if len(dequeue) == 0:
            return True
    return False

def feed_incoming_pipes_wait(data, pipes, rank, target_index):
    smallest_buffer_size = sys.maxint
    for pipe in pipes:
        if pipes[pipe].incoming is True and pipes[pipe].incoming_index == target_index:
#            print data.tostring(), target_index
            buffer_size = pipes[pipe].writeData(data)
            if buffer_size < smallest_buffer_size:
#                print ("Buffer size = %d" %(buffer_size))
                smallest_buffer_size = buffer_size
    
    max_items_in_buffer = 2
    sleep_time = 1     # seconds
    total_sleep_time = 0
    max_sleep_time = 300
    
    if (smallest_buffer_size >= max_items_in_buffer) and (BUFFER_SLEEP is True):
        luigi_logger.debug ("Pausing rank %d, (too greedy) buffer size %d. Continuing checks every %d seconds" %(rank, smallest_buffer_size, sleep_time))
        while True:
            for pipe in pipes:
                if pipes[pipe].incoming is True:
                    buffer_size = pipes[pipe].getQueueSize()
                    if buffer_size < smallest_buffer_size:
                        smallest_buffer_size = buffer_size
            if (smallest_buffer_size <= max_items_in_buffer) or (total_sleep_time >= max_sleep_time):
                break
            else:
                time.sleep(sleep_time)
                total_sleep_time += sleep_time
            